import os  # 用于访问环境变量和操作系统功能
from operator import itemgetter

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDatabaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

api_key = os.getenv("DASHSCOPE_API_KEY")
# 检查API密钥是否存在，若不存在则抛出错误提示
if not api_key:
    raise ValueError("请设置环境变量DASHSCOPE_API_KEY（阿里云百炼API-KEY）")
# 1. 创建大语言模型实例（适配阿里云百炼的DeepSeek模型）
model = ChatOpenAI(
    model="qwen-plus-latest",  # 模型名称，需与阿里云百炼平台支持的模型名一致
    temperature=0,  # 生成内容的随机性（0表示确定性输出，1表示最大随机性）
    max_tokens=None,  # 最大生成 tokens 数（None表示使用模型默认值）
    timeout=None,  # 超时时间（None表示无超时限制）
    max_retries=2,  # 调用失败时的最大重试次数
    api_key=SecretStr(api_key),  # 以安全字符串形式传递API密钥
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 阿里云百炼的OpenAI兼容接口地址
)

# SQLAIChemy 连接MySQL数据库
HOSTNAME = '127.0.0.1'
PORT = '3306'
DATABASE = 'quickstarts'
USERNAME = 'root'
PASSWORD = '492396'
# mysqlclient 驱动URL
MYSQL_URL = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)
# 获得连接
db = SQLDatabase.from_uri(MYSQL_URL)
# 测试连接
# print(db.get_usable_table_names())
# print(db.run('select * from users limit 10'))

# 直接使用大模型和数据库整合
# 初始化生成SQL的Chain
create_sql_chain = create_sql_query_chain(model, db)
# 测试
# resp = create_sql_chain.invoke({'question': '请问users表里面有多少条记录？'})
# print(resp)

# prompt提示词
answer_prompt = PromptTemplate.from_template(
    """给定以下用户问题、SQL语句和SQL执行后的结果，回答用户问题，
    Question: {question}
    SQL Query: {query}
    SQL Result: {result}
    回答:"""
)

#  创建一个执行SQL工具
execute_sql_tool = QuerySQLDatabaseTool(db=db)

# 移除多余的execute_sql_tool调用，并修复工具调用方式
chain = (
    RunnablePassthrough.assign(query=create_sql_chain)  # 生成SQL查询
    | RunnablePassthrough.assign(  # 执行SQL并获取结果
        result=lambda x: execute_sql_tool.invoke(x["query"])  # 使用invoke替代直接调用
    )
    | answer_prompt  # 传入包含question、query、result的字典
    | model
    | StrOutputParser()
)

resp = chain.invoke(input={'question': '请问users表里面有多少条记录？'})
print(resp)
